import numpy as np
from sklearn.datasets import load_iris
from sklearn.tree import DecisionTreeClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import joblib

# 加载鸢尾花数据集
# 鸢尾花数据集（Iris dataset）是一个经典的用于分类任务的数据集，
# 常被用来测试分类算法的性能。这个数据集包含了 150 个样本，
# 分别属于三种不同的鸢尾花（Setosa、Versicolor 和 Virginica），
# 每个样本包含了四个特征：花萼长度、花萼宽度、花瓣长度和花瓣宽度
# 这些特征被用来预测鸢尾花的类别
iris = load_iris()
X = iris.data
y = iris.target

# 将数据集分割为训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 实例化决策树模型
model = DecisionTreeClassifier()

# 模型训练
model.fit(X_train, y_train)

# 进行预测
y_pred = model.predict(X_test)

# 计算准确率
accuracy = accuracy_score(y_test, y_pred)
print(f"准确率: {accuracy}")

# 保存模型
joblib_file = "decision_tree_model.pkl"
joblib.dump(model, joblib_file)
print(f"模型已保存到 {joblib_file}")

# 加载模型
loaded_model = joblib.load(joblib_file)

# 进行预测
new_data = np.array([[5.0, 3.6, 1.4, 0.2]])
prediction = loaded_model.predict(new_data)
print(f"预测值: {iris.target_names[prediction][0]}")

